#include "planner/memo.h"

#include <format>
#include <string>
#include <unordered_set>

#include "planner/group_expression.h"
#include "planner/logical_operators.h"

namespace optimizer {

GroupExpression *Memo::InsertExpression(GroupExpression *gexpr, group_id_t target_group, bool enforced) {
  // If leaf, then just return
  if (gexpr->Contents()->GetOpType() == OperatorType::LeafOperator) {
    const auto &leaf = gexpr->Contents()->CastTo<LeafOperator>();
    CASCADES_ASSERT(target_group == UNDEFINED_GROUP || target_group == leaf.GetOriginGroup(),
                    "target_group does not match the LeafOperator's group");
    gexpr->SetGroupID(leaf.GetOriginGroup());

    // Let the caller delete!
    // Caller needs the origin_group
    return nullptr;
  }

  gexpr->SetGroupID(target_group);
  // Lookup in hash table
  auto it = group_expressions_.find(gexpr);
  if (it != group_expressions_.end()) {
    CASCADES_ASSERT(*gexpr == *(*it), "GroupExpression should be equal");
    delete gexpr;
    return *it;
  }

  group_expressions_.insert(gexpr);

  // New expression, so try to insert into an existing group or
  // create a new group if none specified
  group_id_t group_id;
  if (target_group == UNDEFINED_GROUP) {
    group_id = AddNewGroup(gexpr);
  } else {
    group_id = target_group;
  }

  Group *group = GetGroupByID(group_id);
  group->AddExpression(gexpr, enforced);
  return gexpr;
}

group_id_t Memo::AddNewGroup(GroupExpression *gexpr) {
  auto new_group_id = group_id_t(groups_.size());

  groups_.push_back(new Group(new_group_id));
  return new_group_id;
}

std::string Memo::ToString() const {
  std::string str;
  std::format_to(std::back_inserter(str), "Memo: {} groups\n", groups_.size());
  for (const auto &group : groups_) {
    str += group->ToString() + "\n";
  }
  return str;
}

}  // namespace optimizer
